#!/usr/bin/env python
'''
Title: Pupil: Eye Tracking Software
main.py
Authors: Moritz Kassner & William Patera
Date: July 12, 2011
Notes: Prototype tested with 2 XBox Live Webcams on Macbook Pro OSX 10.6.7 and Ubuntu 10.10 using OpenCV 2.2
'''
import cv
import numpy as np
from init import *
from calibrate import *
from methods import *
from kalman_filter import *


if __name__=="__main__":
	#cam and ROI size vars
	worldcam_size = (640,480)
	eyecam_size = (640,480)
	
	roi_pos = (60,120)
	roi_size = (360,280)
	 
	#interface variables
	start_calibrate = False
	collect_cal_data = False
	stop_calibrate = False
	oneshot = False
	record = False
	#output_track = False
	tick = 0
	draw_blur = False
	draw = False
 
	start_collect_pts = False
	stop_collect_pts = False
	collect_pts = False
	shot = False
	
	#pupil spot default threshold
	p_threshold = 50

	#inital mapping coefficiants
	coefs = None
	pt_cloud = [[None,None]]
	#some internal setup
	roi = roi_pos + roi_size
	eyecam_center = (eyecam_size[0]/2,eyecam_size[1]/2)
	worldcam_center = (worldcam_size[0]/2,worldcam_size[1]/2)
	
	kalman = create_kalman_preset()
	
	
	#initialize capture
	eyecam = capture_init(0,eyecam_size[0],eyecam_size[1],30)
	worldcam =	capture_init(1,worldcam_size[0],worldcam_size[1],30)

	#create windows for each of the images and GUI sliders
	cv.NamedWindow("Eye", 1)
	cv.MoveWindow("Eye", 650, 0)
	cv.NamedWindow("Eye_Filtered", 1)
	cv.MoveWindow("Eye_Filtered", 742, 532)
	cv.NamedWindow("Eye_Gray", 1)
	cv.MoveWindow("Eye_Gray", 0, 532)
	cv.NamedWindow("Eye_Smooth", 1)
	cv.MoveWindow("Eye_Smooth", 372, 532)
	cv.NamedWindow("World", 1)
	cv.MoveWindow("World", 0, 0)
	cv.NamedWindow("Slider",cv.CV_WINDOW_NORMAL)
	cv.MoveWindow("Slider", 1110, 552)
	cv.ResizeWindow("Slider", 200, 60) 
	cv.CreateTrackbar("Slider", "Slider", p_threshold, 140, UpdateVar)
	
	#print user control message to the terminal
	gui_control_msg()
	#video_writer = cv.CreateVideoWriter("out.avi", cv.CV_FOURCC('I', '4', '2', '0'), 30, worldcam_size, 1) 
	while True:
		#grab new frames
		eye = cv.QueryFrame(eyecam)
		world = cv.QueryFrame(worldcam)
		
		#take a subregion of the picture aka Region of interest
		cv.SetImageROI(eye,roi)
		#get the filtered image of the puil and its moment
		eye_filtered, pupil_pos = extract_darkspot_ellipse_center(eye,p_threshold)
		cv.ResetImageROI(eye)
		 
		
		if pupil_pos: #if there is a pupil dedected update the position
			
			#draw the mesured pupil center
			tx,ty = pupil_pos
			eye_pos = (roi_pos[0]+tx, roi_pos[1]+ty)
			x = int(eye_pos[0])
			y = int(eye_pos[1])
			eye_pos = (x, y)
			cv.Circle(eye, eye_pos, 3, cv.CV_RGB(0, 255, 0), thickness=1, lineType=cv.CV_AA, shift=0) 

					
			eye_norm_x = (tx-roi_size[0]/2)/(roi_size[0]/2.)
			eye_norm_y = (ty-roi_size[1]/2)/(roi_size[1]/2.)
			kalman_pos = kalman.update((eye_norm_x ,eye_norm_y))

			# mapping eye pos onto the world
			world_norm_pos = map_vector((kalman_pos),coefs)
			world_screen_pos = [p*float(wc)+wc for p,wc in zip(world_norm_pos,worldcam_center)]
			world_screen_pos = (world_screen_pos[0],world_screen_pos[1])
			
			#un offsetting and unnormalizing the kalman pos to render on the eye
			eye_pos = tuple([p*(rs/2.)+rs/2.+rp for p,rs,rp in zip(kalman_pos,roi_size,roi_pos)])
		
		#two stage calibration routine
		if start_calibrate:
			pt_cloud = [] 
			collect_cal_data = True
			start_calibrate = False
			oneshot = True
		
		if collect_cal_data:
			world,mean = chessboard(world)
		   
			
			if pupil_pos and mean:
				world_obj_x = (mean[0]-worldcam_center[0])/float(worldcam_center[0])
				world_obj_y = (mean[1]-worldcam_center[1])/float(worldcam_center[1])
				#print eye_norm_x,eye_norm_y,world_obj_x,world_obj_y
				pt_cloud.append([eye_norm_x,eye_norm_y,world_obj_x,world_obj_y])
		else: 
			world_filtered = world
			
		if stop_calibrate and oneshot:
			collect_cal_data = False
			stop_calibrate = False
			oneshot = False
			if pt_cloud:
				coefs = calibrate_poly(pt_cloud)
				print coefs
		
		### Collect points for heat map drawing and save to .npy file
		if start_collect_pts:
			track_pts = []
			collect_pts = True
			shot = True
			start_collect_pts = False
		
		if collect_pts:
			track_pts.append(world_screen_pos)
			#track_pts.append(eye_pos) #for testing only
			
		
		if stop_collect_pts and shot:
			start_collect_pts = False
			collect_pts = False
			shot = False
			if track_pts:
				np.save("track_points.npy", track_pts)
				np.savetxt("track_pts.txt", track_pts)
				print "Points saved as .npy file"
			else:
				print "No points saved to file, no points collected"
		
		
		#draw shapes and video on screen
		if pupil_pos and draw_blur:
			#optional drawing type - makes Foveal area in focus, peripheral area blurred
			draw = True
			mask = cv.CreateImage(worldcam_size, 8, 1)
			mask_inv = cv.CreateImage(worldcam_size, 8, 1)
			blur_world = cv.CreateImage(worldcam_size, 8, 3)
			wx = int(world_screen_pos[0])
			wy = int(world_screen_pos[1])
			cv.Circle(mask, (wx, wy), 40, cv.CV_RGB(255, 255, 255), thickness=-1, lineType=cv.CV_AA, shift=0)
			#cv.Not(mask, mask_inv)
			#cv.Smooth(world, blur_world, smoothtype=cv.CV_BLUR, param1=5, param2=5)
			#cv.Copy(blur_world,world,mask_inv)
			
		if pupil_pos and not draw_blur:
			#draws white circle in eye cam and green circle in world cam
			#For some reason the Circle function is not working.  The second paramater.  What is the problem
			x = int(eye_pos[0].item(0))
			y = int(eye_pos[1].item(0))
			#eye_pos = (x,y)
			wx = int(world_screen_pos[0])
			wy = int(world_screen_pos[1])
			#world_screen_pos = (wx,wy)			
			cv.Circle(eye,(x,y), 3, cv.CV_RGB(255, 255, 255), thickness=1, lineType=cv.CV_AA, shift=0) 
			cv.Circle(world, (wx,wy), 12, cv.CV_RGB(0, 255, 0), thickness=1, lineType=cv.CV_AA, shift=0)
			
		cv.ShowImage("Eye", eye)
		cv.ShowImage("World", world)
		cv.ShowImage("Eye_Filtered", eye_filtered)
		
		#record video
		if record:
			#print cv.WriteFrame(video_writer,world)
			#add the tick part!!!!
			cv.SaveImage("out/0"+str(tick)+".PNG", world) 
			tick += 1
		
		#poll trackbar
		p_threshold = cv.GetTrackbarPos("Slider","Slider")
		#keyborad interrupt routines
		c = cv.WaitKey(7) % 0x100
		if c == 27:
			break
		if c == 32: #space bar
			if not collect_cal_data:
				start_calibrate = True
				print "START CALIBRATION Recording"
			elif collect_cal_data and oneshot:
				stop_calibrate = True
				print "STOP CALIBRATION Recording"

		if c == ord("r"):
			if record:
				record = False
				print "Record", record
			else:
				record = True
				print "Record", record
		if c == ord("d"):
			print "Draw", draw
			if not draw:
				draw_blur = True
				print "START DRAWING Foveal Blur"
			elif draw and draw_blur:
				draw_blur = False
				draw = False
				print "STOP DRAWING Foveal Blur"
		if c == ord("t"):
			if not collect_pts:
				start_collect_pts = True
				print "COLLECTING Gaze Points" 
			elif collect_pts and shot:
				stop_collect_pts = True
				print "SAVING Gaze Points to file"
				
